import logging
from sqlalchemy import inspect

class SchemaReader:
    def __init__(self, db_connector):
        self.db_connector = db_connector
        self.logger = logging.getLogger(__name__)
        self.tables_info = {}
    
    def read_all_tables(self):        
        try:
            inspector = inspect(self.db_connector.engine)
            tables = inspector.get_table_names()
            
            for table_name in tables:
                self.tables_info[table_name] = self._get_table_info(inspector, table_name)
            
            return self.tables_info
        except Exception as e:
            self.logger.error(f"Error reading database schema: {str(e)}")
            raise
    
    def _get_table_info(self, inspector, table_name):
        columns = inspector.get_columns(table_name)
        pk_constraint = inspector.get_pk_constraint(table_name)
        primary_keys = pk_constraint.get('constrained_columns', []) if pk_constraint else []
        foreign_keys = inspector.get_foreign_keys(table_name)
        
        
        return {
            "name": table_name,
            "columns": columns,
            "primary_keys": primary_keys,
            "foreign_keys": foreign_keys
        }
    
    
    def get_table_info(self, table_name):
        if not self.tables_info:
            self.read_all_tables()
        
        return self.tables_info.get(table_name)
    
    def get_schema_for_llm(self):
        if not self.tables_info:
            self.read_all_tables()
        
        schema_text = "Database Schema Information:\n\n"
        
        for table_name, table_info in self.tables_info.items():
            schema_text += f"Table: {table_name}\n"
            schema_text += "Columns:\n"
            
            for column in table_info['columns']:
                column_name = column['name']
                column_type = str(column['type'])
                is_nullable = "Nullable" if column.get('nullable', True) else "Not Null"
                is_pk = "Primary Key" if column_name in table_info['primary_keys'] else ""
                
                schema_text += f"  - {column_name} ({column_type}), {is_nullable} {is_pk}\n"
            
            if table_info['foreign_keys']:
                schema_text += "Foreign Keys:\n"
                for fk in table_info['foreign_keys']:
                    schema_text += f"  - {', '.join(fk['constrained_columns'])} -> {fk['referred_table']}.{', '.join(fk['referred_columns'])}\n"
            
            schema_text += "\n"
        
        return schema_text